from openai import OpenAI
import json
from BaiduSearch import baidu_search
import re
client = OpenAI(
  base_url = 'http://localhost:11434/v1/',
  api_key = 'zyy',
)
system_prompt = """
你在“思考→行动→观察→回答”的循环中运行。
在循环结束时，你输出答案。

- 用 思考 描述你对问题的想法。
- 用 行动 执行你可用的一个行动。
- 观察 是执行那些行动后的结果。
- 回答 是你分析 观察 得出的结论。

可用的行动如下：
calculate: 例如 calculate: 4 * 7 / 3 进行计算并返回数值——使用 Python，必要时请使用浮点数语法。

百度搜索: 例如 百度搜索: 中国的首都是什么？ 通过查阅 百度搜索API 返回真实信息。

如果有机会，总是使用 百度搜索 查询。

示例会话：

问题: 中国的首都是什么？ 
思考: 我应该在 百度搜索 上查询。 
行动: 百度搜索: 中国的首都是什么？ 
PAUSE
直接返回结果，等待系统对你的行动进行工具的调用。

随后你将收到：

观察: 中国是一个国家。首都是北京。 
思考: 我认为已经找到了答案。

你随后输出：

回答: 中国的首都是北京。

另一个示例会话：

问题: 地球质量乘以 2 是多少？ 
思考: 我需要在 百度搜索 上查找地球质量。 
行动: 百度搜索: 地球质量 
PAUSE
直接返回结果，等待系统对你的行动进行工具的调用。

随后你将收到：

观察: 地球质量为 1.1944×10e25。

思考: 我需要把它乘以 2。 
行动: calculate: 5.972e24 * 2      
  PAUSE
  直接返回结果，等待系统对你的行动进行工具的调用。

随后你将收到：

观察: 1.1944×10e25。

如果你已得到答案，请将其作为 回答 输出。

回答: 地球质量乘以 2 的结果是 1.1944×10e25。



"""

def calculate(question):
    try:
        result = eval(question)
        return f"{result}"
    except Exception as e:
        return f"计算出错：{e}"
    
available_actions = {
    '百度搜索':baidu_search,
    '科学计算':calculate,
    'calculate':calculate
}

class ChatBot:
    def __init__(self,system=''):
        self.system = system
        self.messages = []
        if self.system:
            self.messages.append({'role':'system','content':system})
    def execute(self):
        client= OpenAI(
          base_url = 'http://localhost:11434/v1/',
          api_key = 'zyy',
        )
        completion = client.chat.completions.create(
          model ='qwen3:4b',
          messages=self.messages
        
        )
        return completion.choices[0].message.content
    def __call__(self,message):
        self.messages.append({'role':'user','content':message})
        answer = self.execute()
        self.messages.append({'role':'assistant','content':answer})
        return answer

action_re = re.compile('^行动: (\w+): (.*)$')
def AgentExecute(question,max_turns=5):
    i=0
    bot = ChatBot(system=system_prompt)
    #通过next_prompt 标识子任务的阶段性输入
    next_prompt = question
    while i<max_turns:
        i+=1
        print(f"\n第 {i} 轮对话")
        #调用ChatBot 的__call__方法
        
        result = bot(next_prompt)
        print(f'result:{result}')
        actions = [action_re.match(a) for a in result.split('\n') if action_re.match(a)]
        if actions:
            # 提取调用的工具名和工具所需的入参
            action, action_input = actions[0].groups()
            if action not in available_actions:
                raise Exception("未知行动: {}: {}".format(action, action_input))
            print(f"正在执行: {action} {action_input}")
            observation = available_actions[action](action_input)
            print(f"观察: {observation}")
            next_prompt = f"观察: {observation}"
        else:
            return bot.messages
bot = ChatBot(system=system_prompt)
question ='韩国最高楼的高度的二次方是多少？'
AgentExecute(question)



